from llama_index.core import SimpleDirectoryReader, PromptHelper, StorageContext, ServiceContext, GPTVectorStoreIndex, load_index_from_storage
from llama_index.core import LLMPredictor
from langchain.chat_models import ChatOpenAI
import gradio as gr
import sys
import os
import openai


openai.api_base = "API URL"
API_TOKEN = 'YOUR TOKEN'
os.environ["OPENAI_API_KEY"] = API_TOKEN
openai.api_key = API_TOKEN


# 创建服务上下文
def create_service_context():
    # 最大输入大小
    max_input_size = 4096
    # 输出数量
    num_outputs = 3072
    # 块重叠的比例
    chunk_overlap_ratio = 0.1
    # 块大小限制
    chunk_size_limit = 600

    # - PromptHelper主要负责生成适当的提示语句，使模型能够理解并有效地响应我们的查询。
    # - LLM Predictor则负责调用gpt语言模型来生成我们期望的回答。
    prompt_helper = PromptHelper(max_input_size, num_outputs, chunk_overlap_ratio, chunk_size_limit=chunk_size_limit)
    llm_predictor = LLMPredictor(llm=ChatOpenAI(temperature=0.5, model_name="gpt-3.5-turbo", max_tokens=num_outputs))
    # 使用默认值构造 service_context
    service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor, prompt_helper=prompt_helper)
    return service_context


# llama index+langchain 数据提取转换为矢量数据，并建立索引数据
def data_ingestion_indexing(directory_path):
    # 从指定的目录路径加载数据。这里假设SimpleDirectoryReader是一个已经定义好的类，它的load_data方法用于从指定目录加载数据
    documents = SimpleDirectoryReader(directory_path).load_data()
    # 当第一次构建索引时，使用GPTVectorStoreIndex的from_documents方法。这个方法需要一个文档列表和一个服务上下文作为参数，
    index = GPTVectorStoreIndex.from_documents(
        documents, service_context=create_service_context()
    )
    # 将索引持久化到磁盘上，这里默认存储在"storage"文件夹中。这里假设index对象有一个storage_context属性，它有一个persist方法用于将索引持久化到磁盘上
    index.storage_context.persist()
    # 返回构建的索引
    return index


# 数据查询
def data_querying(input_text):
    storage_context = StorageContext.from_defaults(persist_dir="./storage")
    # 从存储中加载索引。load_index_from_storage是一个函数，它接受存储上下文和服务上下文作为参数，返回加载的索引
    index = load_index_from_storage(storage_context, service_context=create_service_context())
    # 使用输入文本查询索引。as_query_engine方法将索引转换为查询引擎，query方法执行查询，并接受输入文本作为参数
    response = index.as_query_engine().query(input_text)
    # 返回查询的响应。response属性包含了查询的结果
    return response.response


iface = gr.Interface(fn=data_querying,
                     inputs=gr.components.Textbox(lines=7, label="运维大模型知识库"),
                     outputs="text",
                     title="GPT碧茂知识库")
# 传递目录
index = data_ingestion_indexing("data")
iface.launch(share=False)